

#include "suffix_tree.h"
#include "util.h"

SuffixTree::SuffixTree(TREE_TYPE setType)
{
	minMumWordLength = 2;
	root = new SuffixTreeNode;
	type = setType;
}

SuffixTree::~SuffixTree()
{

	root->free();
	delete root;
}

void SuffixTree::insert(Sentence * sen, wchar_t *wstring)
{
	appendAsChildOrStringTail(root,sen,wstring,wstring);
}


void SuffixTree::insert(SuffixTreeNode * cur, Sentence *s, wchar_t *begin, wchar_t* wstring)
{
	cur->counter++;
	if(type == SuffixTree::RIGHT_SUFFIX_TREE && *(++wstring) != 0
		|| type == SuffixTree::LEFT_SUFFIX_TREE && *(--wstring) != 0)
	{
		appendAsChildOrStringTail(cur,s,begin,wstring);
	}else
	{
		cur->endWithNullCounter++;
	}
}

void SuffixTree::appendAsChildOrStringTail(SuffixTreeNode * cur, Sentence *s, wchar_t *begin, wchar_t* wstring)
{
	int i = 0;
	int temp = cur->children_size;
	for(i = 0; i < temp; ++i)
	{
		//declared as a pointer
		SuffixTreeNode* c = &cur->nodeList[i];
		if(c->state == *wstring)
		{
			//record the start and end pointer of this prefix string in its original sentence
			addSentenceSection(c, s,begin, wstring);

			insert(c,s,begin,wstring);
			break;
		}	
	}
	if(i == temp)
	{
		WordBeginEnd * cursor;
		
		for(cursor=cur->tailHead->next; cursor != 0; cursor = cursor->next)
		{
			if(*wstring == *cursor->end)
			{
				
				// we need to add append new node at current node
				// and remove this suffix at current node , insert them at new node 
				SuffixTreeNode *newNode = cur->addChildNode(*wstring);

				Sentence *s2 = cursor->sen;
				wchar_t *b2 = cursor->begin;
				wchar_t *e2 = cursor->end;
				
				// remove the old string tail
				cur->removeStringTail(cursor);

				addSentenceSection(newNode,s,begin,wstring);
				addSentenceSection(newNode,s2,b2,e2);

				insert(newNode, s, begin, wstring);
				insert(newNode, s2,b2,e2);

				break;
			}
		}
		if( cursor == 0)
		{
			//just append as a suffix at current node
			cur->addStringTail(s,begin,wstring);
		}
	}
}

void SuffixTree::addSentenceSection(SuffixTreeNode *node,Sentence *s, wchar_t *begin, wchar_t *end)
{
	if(type == SuffixTree::RIGHT_SUFFIX_TREE)
	{
		node->addSentenceSection(s,begin,end);
	}else
	{
		node->addSentenceSection(s,end,begin);
	}
}

void SuffixTree::collectWordContext(WordContext *& list, int& size)
{

	// we'll only collect the context of word whose depth(length) is larger than 2.
	int listSize = countNodes(root,-1,2);
	list = new WordContext[listSize];
	
	size = 0;		// in this loop it is sure that size will increment to listSize.
	for(int i = 0; i < root->children_size; i ++)
	{
		collectWordContext(&root->nodeList[i],		//current root
							0,						//parent depth 
							0,						//prefix
							list,					//WordContext list
							size);					//index
	}
}

void SuffixTree::collectWordContext(SuffixTreeNode *node, int parent_depth, wchar_t* prefix, WordContext* list,int& index)
{
	//TODO: we may never collect word of length less than 2.
	
	int depth = parent_depth+1;
	
	if(depth < minMumWordLength)
	{
		wchar_t * temp = new wchar_t[depth];
		appendToPrefix(prefix,node->state,temp,parent_depth);

		//travel to all node recursively
		for( int i = 0; i < node->children_size; ++i)
		{
			collectWordContext(&node->nodeList[i],depth, temp,list,index);
		}
		delete [] temp;
	}
	else
	{
		WordContext * ctx = &list[index++];

		//initial length
		ctx->length = depth;
		
		//initial str
		ctx->str = new wchar_t[depth];
		
		appendToPrefix(prefix,node->state,ctx->str, parent_depth);

		//remember current node reference
		ctx->treeNode = node;

		//travel to all node recursively
		for( int i = 0; i < node->children_size; ++i)
		{
			collectWordContext(&node->nodeList[i],depth,ctx->str,list,index);
		}
	}
}

void SuffixTree::appendToPrefix(wchar_t * prefix, wchar_t state, wchar_t *newResult, int prefix_length)
{
	if(type == SuffixTree::RIGHT_SUFFIX_TREE)
	{
		for(int i = 0; i < prefix_length; ++i)
		{
			newResult[i] = prefix[i];
		}
		newResult[prefix_length] = state;
	}
	else
	{
		newResult[0] = state;
		for(int i = 0; i < prefix_length; ++i)
		{
			newResult[i+1] = prefix[i];
		}
	}

}


void SuffixTree::travel()
{
	printf("Total %d nodes.\n",countNodes(root,-1,0));
	printf("%d Words have been inserted.\n",countWordInserted(root));
	
	printf("-level:%d, suffix:%d children:%d\n",0,root->stringTailLength,root->children_size);
	
	for(WordBeginEnd *cur = root->tailHead->next; cur != 0; cur=cur->next)
	{
		printf("  ");
		if(type == SuffixTree::RIGHT_SUFFIX_TREE)
			print_wchar(cur->end);
		else
			print_wchar_reverse(cur->end);
		printf("\n");
	}

	for(int i = 0; i < root->children_size; i++)
	{
		travel(&root->nodeList[i],0,1);
	}
}

void SuffixTree::travel(SuffixTreeNode *node,wchar_t *prefix, int depth)
{
	wchar_t *str = new wchar_t[depth+1];
	wchar_t *next_pre = new wchar_t[depth];
	for(int i = 0; i <depth-1; ++i)
	{
		str[i] = prefix[i];
		next_pre[i] = prefix[i];
	}

	str[depth-1] = node->state;
	next_pre[depth-1] = node->state;
	str[depth] = 0;

	for(int i = 0; i < depth; i++)
		printf("  ");
	printf("prefix: ");
	print_wchar(str);
	
	for(int i = 0; i < depth; i++)
		printf("  ");
	printf("-level:%d, suffix:%d children:%d count:%d \n",depth,node->stringTailLength,node->children_size,node->counter);
	
	for(WordBeginEnd *cur = root->tailHead->next; cur != 0; cur=cur->next)
	{
		for(int i = 0; i < depth; i++) printf("  ");
		printf(" ");
		if(type == SuffixTree::RIGHT_SUFFIX_TREE)
			print_wchar(cur->end);
		else
			print_wchar_reverse(cur->end);
		printf("\n");
	}

	for(int i = 0; i < node->children_size; i++)
	{

		travel(&node->nodeList[i], next_pre,depth+1);
		
	}
	delete [] str;
	delete [] next_pre;

}

//------------------------------------------------------
int countNodes(SuffixTreeNode *root,int parent_depth,int depth)
{
	int counter = 0;
	int curdepth = parent_depth+1;

	if( curdepth >= depth)
	{
		counter = 1;
	}

	for(int i = 0; i < root->children_size; ++i)
	{
		counter += countNodes(&root->nodeList[i], curdepth,depth);
	}
	return counter;                                                                                                                                                                                                                                                                                           ;
}

int countWordInserted(SuffixTreeNode *root)
{
	int count=root->stringTailLength + root->endWithNullCounter;
	for(int i = 0;i < root->children_size;++i)
	{
		count += countWordInserted(&root->nodeList[i]);
	}
	return count;

}
